Early review (WIP) -- fastgen DMD2 model integration

Thanks for the PR. The overall approach of subclassing the existing Wan 2.2 T2V / I2V pipelines for DMD2-distilled 4-step models is reasonable. Here are observations for the current state:

1. Heavy code duplication between T2V and I2V DMD2 classes

WanT2VDMD2Pipeline and WanI2VDMD2Pipeline share identical implementations for:

• __init__ (scheduler replacement)
• _verify_dmd2_request (entire method, ~40 lines)
• forward timestep-patching logic (monkey-patch + try/finally)
• Class constants (GUIDANCE_SCALE, NUM_INFERENCE_STEPS, DMD2_TIMESTEPS)

Please extract this into a mixin (e.g. DMD2PipelineMixin) that both classes inherit from. This avoids the maintenance burden of keeping two copies in sync and makes it trivial to add future DMD2 variants (e.g. VACE).

2. Monkey-patching scheduler.set_timesteps is fragile

In forward(), you replace self.scheduler.set_timesteps with a closure and restore it in a finally block. This approach has problems:

• Not thread-safe / async-safe: if two requests are in flight concurrently on the same pipeline instance, they will race on the scheduler's method.
• Fragile against parent refactors: if the parent ever calls set_timesteps differently (e.g. via a local reference), the patch silently breaks.

Consider instead:

• Overriding the parent's denoising loop method directly, or
• Calling self.scheduler.set_timesteps(timesteps=self.DMD2_TIMESTEPS, device=...) once after the parent calls it (by wrapping only the relevant section), or
• Using FlowMatchEulerDiscreteScheduler's built-in timesteps parameter at set_timesteps time and overriding num_inference_steps in the super call.

3. NUM_INFERENCE_STEPS class constant is declared but never used

You define NUM_INFERENCE_STEPS = 4 on both classes but never reference it. The forward() signature defaults to num_inference_steps: int = 4 as a separate literal. Either use the constant in the default or remove it to avoid confusion.

4. _verify_dmd2_request mutates the request in-place silently

The method modifies req.sampling_params and req.prompts in-place. While there are log warnings, consider:

• Documenting in the method docstring that the request is mutated (not just verified -- the name _verify_* is misleading for a method that modifies data). A name like _sanitize_dmd2_request would be more accurate.
• Returning the modified request or a copy, so callers see that mutation happens. In-place mutation of shared request objects can cause subtle bugs if the same request is reused.

5. p.get("negative_prompt") is falsy for empty strings

In _verify_dmd2_request:

if isinstance(p, dict) and p.get("negative_prompt"):

This will not strip negative_prompt if its value is "" (empty string). If the intent is to remove the key entirely when present, use "negative_prompt" in p instead.

6. Test helper uses object.__new__ to skip __init__ -- brittle

In test_wan_dmd2_request_sanitization.py, _make_pipeline does:

pipeline = object.__new__(cls)
torch.nn.Module.__init__(pipeline)

This creates a pipeline without running any __init__, so _verify_dmd2_request works only because it doesn't touch any instance state set in __init__. If the method ever accesses self.GUIDANCE_SCALE (which it already does), this works only because it's a class attribute. This is fragile -- consider using the mock-based approach from the scheduler test file instead, which is more robust.

7. Missing import in scheduler test

In test_wan_dmd2_scheduler.py, line 40:

from vllm_omni.diffusion.request import OmniDiffusionRequest, OmniDiffusionSamplingParams

This import is placed in the middle of the file (after fixtures, before test functions). Move it to the top with the other imports for consistency.

8. Unused import in scheduler test

import inspect on line 2 of test_wan_dmd2_scheduler.py is never used. Remove it.

9. Registry entries look correct

The additions to registry.py (pipeline registry, pre/post-process function mappings) look correct and consistent with existing entries. The DMD2 pipelines appropriately reuse the same pre/post-process functions as their parent pipelines.

10. PR description is empty

Please fill in the Purpose, Test Plan, and Test Result sections. Even for WIP, a brief description of the DMD2 distillation approach and which models this targets (model hub IDs) would help reviewers.

Overall, the core idea is sound. The main actionable items before this is ready for merge are: (1) extract the duplicated DMD2 logic into a mixin, (2) find a less fragile approach than monkey-patching set_timesteps, and (3) minor cleanup items above.